import os
import scipy.ndimage as ndimage
import SimpleITK as sitk
from PIL import Image
import numpy as np
import pydicom as dicom
'''
CHAOS的标签
肝：63
右肾（right kidney）：126
左肾（left kidney）：189
脾（spleen）：252
63,126,189,252
'''
# 获取dicom文件的pixel_array
def get_pixel_array(dicom_file_path):
    dcm_ori = dicom.dcmread(dicom_file_path)
    arr = dcm_ori.pixel_array
    return arr
def getLabel(label_array,include_label=[0,63,126,189,252]):
    '''
    label_array：标注数据
    include_label：要筛选出来的标注list，例如[0,1,2,3,6],筛选出四个器官
    '''
    '''
    CHAOS的标签
    肝：63,4
    右肾（right kidney）：126,2
    左肾（left kidney）：189,3
    脾（spleen）：252,1
    63,126,189,252
    '''
    k = np.isin(label_array, np.array(include_label))
    label_array[np.logical_not(k)] = 0
    label_array[label_array == 63] = 4
    label_array[label_array == 126] = 2
    label_array[label_array == 189] = 3
    label_array[label_array == 252] = 1
    return label_array
def processAndSave(data_dir=r"D:\AllData\CHAOS\Train_Sets\Train_Sets\MR",
                   save_dir=r"D:\AllData\CHAOS\afterProcess"):
    if not os.path.exists(os.path.join(save_dir,"volume")):
        os.makedirs(os.path.join(save_dir,"volume"))
    if not os.path.exists(os.path.join(save_dir,"segmentation")):
        os.makedirs(os.path.join(save_dir,"segmentation"))
    number = 0

    subject_list=os.listdir(data_dir)
    subject_list.sort()
    for index,example in enumerate(subject_list):
        #单独对每一个3D数据进行z-score归一化
        image_array_list = []
        label_array_list = []
        image_name_list=[]
        example_ground_dir=os.path.join(data_dir,example,"T2SPIR/Ground")
        for image_name in sorted(os.listdir(example_ground_dir)):
            image_name_list.append(image_name.split('.')[0])
            label_path=os.path.join(example_ground_dir,image_name)
            image_path=label_path.replace('Ground','DICOM_anon').replace('png','dcm')
            last_shape = np.array([256, 256])#要统一的维度
            # -------------------处理标注数据---------------------#
            label = Image.open(label_path).convert("L")
            #将脾的标签设为1，其余器官和背景设为0
            label_array=np.array(label)
            # label_array[label_array!=252]=0
            # label_array[label_array == 252] = 1
            set1=set(label_array.flatten())


            label_array = getLabel(label_array)
            set2 = set(label_array.flatten())


            label_dimension_adjustment = 1 / (label_array.shape / last_shape)
            label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)#只会出现原来出现过的数值
            label_array_list.append(label_array)
            # -------------------处理图像数据---------------------#
            image_array=get_pixel_array(image_path)
            image_dimension_adjustment = 1 / (image_array.shape / last_shape)
            image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)#双线性插值
            image_array_list.append(image_array)
        #以整个3D图像进行z-score,并拼接image和label的每个2D切片，[image, mask] 。并保存
        image_array_list=np.array(image_array_list)
        #image_array_list=(image_array_list-image_array_list.mean())/image_array_list.std()
        label_array_list=np.array(label_array_list)
        label_nii=sitk.GetImageFromArray(label_array_list)
        image_nii=sitk.GetImageFromArray(image_array_list)
        sitk.WriteImage(image_nii,os.path.join(save_dir,'volume',example+'.nii.gz'))
        sitk.WriteImage(label_nii, os.path.join(save_dir, 'segmentation', example + '.nii.gz'))
        print(index)
    print(number)

if __name__ == '__main__':
    processAndSave(data_dir=r"D:\AllData\CHAOS\Train_Sets\Train_Sets\MR",
                   save_dir=r"D:\AllData\CHAOS\Train_Sets\3Dnii")
